from __future__ import annotations
from .codeql_metadata import *


def make_graph_index(nodes:tuple,node_locations:tuple[NodeLocation],node_ids:tuple):
    line2nodes={}
    id2nodes={}
    node2index={}
    
    for i,e in enumerate(zip(nodes,node_locations,node_ids)):
        node,loc,id=e
        row=loc.start_row
        if row not in line2nodes:
            line2nodes[row]=[]
        line2nodes[row].append(node)
        id2nodes[id]=e
        node2index[e]=i

    return line2nodes,id2nodes,node2index

class AstTree:
    def __init__(self,ast_nodes:list[AstNode],ast_edges:list[AstEdge]) -> None:
        self.ast_nodes=ast_nodes
        self.ast_edges=ast_edges
        self.ast_nodes.sort(key=lambda x:x.location.key)

        self.line2nodes,self.id2nodes,self.node2index=make_graph_index(
            self.ast_nodes,
            map(lambda x:x.location,self.ast_nodes),
            map(lambda x:x.id,self.ast_nodes)
        )

    def get_node_by_loc(self,node_loc:NodeLocation):
        start_row=node_loc.start_row
        start_col=node_loc.start_column
        end_col=node_loc.end_column
        
        candiate_nodes=self.line2nodes.get(start_row,[])
        
        for e in candiate_nodes:
            e_start_col=e.location.start_column
            e_end_col=e.location.end_column
            if e_end_col==end_col and e_start_col==start_col:
                return e
        
        return None
    
    def select_node_by_loc(self,start_row,start_col,end_row,end_col):
        result=[]
        multi_line_flag=start_row!=end_row
        
        for line_id in range(start_row,end_row+1):
            tokens=self.line2nodes[line_id]
            if multi_line_flag:
                if line_id==start_row:
                    result+=[e for e in tokens if e.start[1]>=start_col]
                elif line_id==end_row:
                    result+=[e for e in tokens if e.start[1]<=end_col]
                else:
                    result+=tokens
                    
            else:
                result+=[e for e in tokens if e.start[1]>=start_col and e.start[1]<=end_col]

        return result

    def iter_edges(self)->tuple[AstNode,AstNode]:
        for edge in self.ast_edges:
            yield self.id2nodes[edge.start_id],self.id2nodes[edge.end_id]

    def indexof(self,ast_node):
        return self.node2index[ast_node]

    @property
    def file_path(self):
        if len(self.ast_nodes)>0:
            return self.ast_nodes[0].location.file_path
        else:
            return None
        
class TokenSequence:
    def __init__(self,tokens:list[TokenNode]) -> None:
        self.tokens=tokens
        self.line2tokens={}
        self.tokens2index={}
        
        for i,e in enumerate(self.tokens):
            row,_=e.start
            if row not in self.line2tokens:
                self.line2tokens[row]=[]
            self.line2tokens[row].append(e)

            self.tokens2index[e]=i

        # for line_id in self.line2tokens:
        #     self.line2tokens[line_id].sort(key=lambda x:x.start[1])
        
    def get_token_by_loc(self,node_loc:NodeLocation):
        start_row=node_loc.start_row
        start_col=node_loc.start_column-1#tokenize生成的column从0开始,为了对齐，所以-1
        end_col=node_loc.end_column-1
        
        candiate_nodes=self.line2tokens.get(start_row,[])
        
        for e in candiate_nodes:
            _,e_start_col=e.start
            _,e_end_col=e.end
            if e_end_col==end_col and e_start_col==start_col:
                return e
        
        return None
    
    def select_token_by_loc(self,start_row,start_col,end_row,end_col):
        result=[]
        start_col-=1#tokenize生成的column从0开始,为了对齐，所以-1
        end_col-=1#tokenize生成的column从0开始,为了对齐，所以-1
        multi_line_flag=start_row!=end_row
        
        for line_id in range(start_row,end_row+1):
            tokens=self.line2tokens[line_id]
            if multi_line_flag:
                if line_id==start_row:
                    result+=[e for e in tokens if e.start[1]>=start_col]
                elif line_id==end_row:
                    result+=[e for e in tokens if e.start[1]<=end_col]
                else:
                    result+=tokens
                    
            else:
                result+=[e for e in tokens if e.start[1]>=start_col and e.start[1]<=end_col]

        return result
        
    def indexof(self,token:TokenNode):
        return self.tokens2index[token]
        
    def raw_tokens(self):
        return map(lambda x:x.string,self.tokens)

class CfgGraph:
    def __init__(self,cfg_nodes:list[CfgNode],cfg_edge:list[CfgEdge]) -> None:
        self.cfg_nodes=cfg_nodes
        self.cfg_edges=cfg_edge
        
        self.cfg_nodes.sort(key=lambda x:x.location.key)
        
        self.line2nodes,self.id2nodes,self.node2index=make_graph_index(
            self.cfg_nodes,
            map(lambda x:x.location,self.cfg_nodes),
            map(lambda x:x.id,self.cfg_nodes)
        )
        
    def indexof(self,node:CfgNode):
        return self.node2index[node]
  
class BasicBlockGraph:
    def __init__(self,basicblock_info:list[BasicBlockInfo],basicblock_edges:list[BasicBlockEdge]) -> None:
        self.basicblock_nodes=basicblock_info
        self.basicblock_edges=basicblock_edges
        
        self.line2nodes,self.id2nodes,self.node2index=make_graph_index(
            self.basicblock_nodes,
            map(lambda x:x.location,self.basicblock_nodes),
            map(lambda x:x.id,self.basicblock_nodes)
        )

class Source:
    def __init__(self,ast_tree:AstTree,token_seq:TokenSequence,basic_blocks:BasicBlockGraph) -> None:
        self.ast_tree=ast_tree
        self.token_seq=token_seq
        self.basic_group_graph=basic_blocks
    
    
    def iter_ast(self):
        for i,e in enumerate(self.ast_tree.ast_nodes):
            yield e,f'{e.location.start_row},{e.location.start_column}'#,self.token_seq.get_token_by_loc(e.location.start_row,e.location.start_column)
            
            
    def iter_token(self):
        for i,e in enumerate(self.token_seq.tokens):
            yield i,e
            
    def iter_basic_block(self):
        for i,e in enumerate(self.basic_group_graph.basicblock_nodes):
            yield i,e
            
    def iter_pair(self):
        for i,e in enumerate(self.ast_tree.ast_nodes):
            
            tokens=self.token_seq.select_token_by_loc(e.location.start_row,e.location.start_column,e.location.end_row,e.location.end_column)
            
            yield e.text,list(map(lambda x:x.string,tokens)) 
            
    def get_ast_mapping_token(self,ast_node,remove_separators=False):
        def is_separator(token:TokenNode):
            if token.string in ['',' ','.','(',')',',',':','[',']','=']:
                return True
            else:
                return False
        
        
        
        tokens=self.token_seq.select_token_by_loc(ast_node.location.start_row,ast_node.location.start_column,ast_node.location.end_row,ast_node.location.end_column)
        
        if remove_separators:
            tokens=[e for e in tokens if not is_separator(e)]
        
        return tokens
    
    def get_basic_block_mapping_token(self,basic_block:BasicBlockInfo):
        loc=basic_block.block_location
        tokens=self.token_seq.select_token_by_loc(loc.start_row,loc.start_column,loc.end_row,loc.end_column)
    
        return tokens
    
    def generate_ast_graph(self):
        for start,end in self.ast_tree.iter_edges():
            token_for_start=self.get_ast_mapping_token(start,True)
            token_for_end=self.get_ast_mapping_token(end,True)
            yield (start.text,token_for_start),(end.text,token_for_end)
    
    def generate_tokens_for_nn(self):
        tokens=list(map(lambda x:x.string,self.token_seq.tokens))
        return tokens
    
    def generate_ast_for_nn(self):
        
        ast_node2token=[(node.node_type,self.get_ast_mapping_token(node)) for node in self.ast_tree.ast_nodes]
        ast_node2token=[(node,list(map(self.token_seq.indexof,tokens))) for node,tokens in ast_node2token]
        
        ast_edges=[(self.ast_tree.indexof(start),self.ast_tree.indexof(end)) for start,end in self.ast_tree.iter_edges()]
        
        return ast_node2token,ast_edges
    
    